# test_nonlinear.py

from mindspore import context
from mindspore.communication.management import init
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
init()

import numpy as np
from mindspore import dataset as ds
from mindspore import nn, Tensor, Model
import time
from mindspore.train.callback import Callback, LossMonitor, ModelCheckpoint, CheckpointConfig
from mindspore.context import ParallelMode
import mindspore as ms

from mindspore.communication.management import init, get_rank, get_group_size

ms.common.set_seed(0)

start_time = time.time()
def get_data(num, a=2.0, b=3.0, c=5.0):
    for _ in range(num):
        x = np.random.uniform(-1.0, 1.0)
        y = np.random.uniform(-1.0, 1.0)
        noise = np.random.normal(0, 0.03)
        z = a * x ** 2 + b * y ** 3 + c + noise
        yield np.array([[x**2], [y**3]],dtype=np.float32).reshape(1,2), np.array([z]).astype(np.float32)

def create_dataset(num_data, batch_size=16, repeat_size=1, shard_id=0, num_shards=2):
    input_data = ds.GeneratorDataset(list(get_data(num_data)), column_names=['xy','z'], num_shards=num_shards, shard_id=shard_id)
    input_data = input_data.batch(batch_size)
    input_data = input_data.repeat(repeat_size)
    return input_data

data_number = 160
batch_number = 16
repeat_number = 20

init("nccl")
rank_id = get_rank()

# 设置GPU个数
device_num = get_group_size()

context.set_auto_parallel_context(device_num=device_num, gradients_mean=False, parallel_mode=ParallelMode.DATA_PARALLEL,
                                  parameter_broadcast=True)
ds_train = create_dataset(data_number, batch_size=batch_number, repeat_size=repeat_number, shard_id=rank_id, num_shards=device_num)
dict_datasets = next(ds_train.create_dict_iterator())

class LinearNet(nn.Cell):
    def __init__(self):
        super(LinearNet, self).__init__()
        self.fc = nn.Dense(2, 1, 0.02, 0.02)

    def construct(self, x):
        x = self.fc(x)
        return x

net = LinearNet()
model_params = net.trainable_params()
print ('Param Shape is: {}'.format(len(model_params)))
for net_param in net.trainable_params():
    print(net_param, net_param.asnumpy())
net_loss = nn.loss.MSELoss()

optim = nn.Momentum(net.trainable_params(), learning_rate=0.01, momentum=0.6)
ckpt_config = CheckpointConfig()
ckpt_callback = ModelCheckpoint(prefix='data_parallel', config=ckpt_config)

model = Model(net, net_loss, optim)

epoch = 50
model.train(epoch, ds_train, callbacks=[ckpt_callback], dataset_sink_mode=True)

for net_param in net.trainable_params():
    print(net_param, net_param.asnumpy())

print ('The total time cost is: {}s'.format(time.time() - start_time))